#!/bin/bash

# 5GC Testbed Shutdown Script
# This script properly shuts down all 5GC components in the correct order

set -e  # Exit on any error

# Configuration
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
}

# Function to check if a service is running
check_service() {
    local service=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "systemctl is-active --quiet $service" 2>/dev/null; then
        return 0
    else
        return 1
    fi
}

# Function to check if a process is running
check_process() {
    local process=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "pgrep -f '$process'" >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to wait for service to stop
wait_for_service_stop() {
    local service=$1
    local vm=$2
    local max_attempts=30
    local attempt=1
    
    log "Waiting for $service on $vm to stop..."
    
    while [ $attempt -le $max_attempts ]; do
        if ! check_service $service $vm; then
            log "✓ $service on $vm has stopped"
            return 0
        fi
        
        echo -n "."
        sleep 2
        ((attempt++))
    done
    
    warn "$service on $vm may not have stopped completely"
    return 1
}

# Function to wait for process to stop
wait_for_process_stop() {
    local process=$1
    local vm=$2
    local max_attempts=30
    local attempt=1
    
    log "Waiting for $process on $vm to stop..."
    
    while [ $attempt -le $max_attempts ]; do
        if ! check_process $process $vm; then
            log "✓ $process on $vm has stopped"
            return 0
        fi
        
        echo -n "."
        sleep 2
        ((attempt++))
    done
    
    warn "$process on $vm may not have stopped completely"
    return 1
}

# Pre-flight checks
log "Performing pre-flight checks..."

# Check VM connectivity
for vm in $VM1 $VM2 $VM3 $VM4; do
    if ping -c 1 -W 5 $vm >/dev/null 2>&1; then
        log "✓ VM $vm is reachable"
    else
        warn "VM $vm is not reachable"
    fi
done

# Check current status
log "Checking current service status..."

# Check what's running
running_services=()
running_processes=()

# Check Open5GS services
if check_service "open5gs-amfd" $VM1; then
    running_services+=("open5gs-amfd")
fi
if check_service "open5gs-smfd" $VM1; then
    running_services+=("open5gs-smfd")
fi
if check_service "open5gs-udmd" $VM1; then
    running_services+=("open5gs-udmd")
fi
if check_service "open5gs-udrd" $VM1; then
    running_services+=("open5gs-udrd")
fi
if check_service "open5gs-pcfd" $VM1; then
    running_services+=("open5gs-pcfd")
fi
if check_service "open5gs-upfd" $VM2; then
    running_services+=("open5gs-upfd")
fi

# Check UERANSIM processes
if check_process "nr-gnb" $VM3; then
    running_processes+=("nr-gnb")
fi
if check_process "nr-ue" $VM4; then
    running_processes+=("nr-ue")
fi

if [ ${#running_services[@]} -eq 0 ] && [ ${#running_processes[@]} -eq 0 ]; then
    log "No 5GC services are currently running"
    exit 0
fi

log "Found running services: ${running_services[*]}"
log "Found running processes: ${running_processes[*]}"

# Confirmation prompt
echo
read -p "Are you sure you want to shut down the 5GC testbed? (y/N): " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    log "Shutdown cancelled"
    exit 0
fi

# Step 1: Stop User Equipment (VM4) - First to disconnect UEs
if check_process "nr-ue" $VM4; then
    log "Stopping UE on VM4..."
    ssh $VM4 "pkill -f nr-ue"
    wait_for_process_stop "nr-ue" $VM4
else
    log "UE is not running on VM4"
fi

# Step 2: Stop Radio Access Network (VM3) - Disconnect gNodeB
if check_process "nr-gnb" $VM3; then
    log "Stopping gNodeB on VM3..."
    ssh $VM3 "pkill -f nr-gnb"
    wait_for_process_stop "nr-gnb" $VM3
else
    log "gNodeB is not running on VM3"
fi

# Step 3: Stop User Plane Function (VM2) - Stop UPF
if check_service "open5gs-upfd" $VM2; then
    log "Stopping UPF on VM2..."
    ssh $VM2 "sudo systemctl stop open5gs-upfd"
    wait_for_service_stop "open5gs-upfd" $VM2
else
    log "UPF is not running on VM2"
fi

# Step 4: Stop Core Network Services (VM1) - Stop in reverse order
log "Stopping core network services on VM1..."

# Stop core services in reverse order
services=("open5gs-amfd" "open5gs-smfd" "open5gs-pcfd" "open5gs-udmd" "open5gs-udrd")

for service in "${services[@]}"; do
    if check_service $service $VM1; then
        log "Stopping $service..."
        ssh $VM1 "sudo systemctl stop $service"
        wait_for_service_stop $service $VM1
        sleep 2
    else
        log "$service is not running"
    fi
done

# Step 5: Stop MongoDB (optional)
read -p "Do you want to stop MongoDB as well? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    if check_service "mongod" $VM1; then
        log "Stopping MongoDB..."
        ssh $VM1 "sudo systemctl stop mongod"
        wait_for_service_stop "mongod" $VM1
    else
        log "MongoDB is not running"
    fi
fi

# Final verification
log "Performing final verification..."

# Check if any services are still running
still_running=()

for service in "${services[@]}" "open5gs-upfd"; do
    if check_service $service $VM1 || check_service $service $VM2; then
        still_running+=($service)
    fi
done

# Check if any processes are still running
if check_process "nr-gnb" $VM3; then
    still_running+=("nr-gnb")
fi
if check_process "nr-ue" $VM4; then
    still_running+=("nr-ue")
fi

if [ ${#still_running[@]} -eq 0 ]; then
    log "✓ All 5GC services have been stopped successfully"
else
    warn "The following services may still be running: ${still_running[*]}"
    
    # Force kill if needed
    read -p "Do you want to force kill remaining processes? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        log "Force killing remaining processes..."
        
        for service in "${still_running[@]}"; do
            case $service in
                "nr-gnb")
                    ssh $VM3 "pkill -9 -f nr-gnb" 2>/dev/null || true
                    ;;
                "nr-ue")
                    ssh $VM4 "pkill -9 -f nr-ue" 2>/dev/null || true
                    ;;
                *)
                    ssh $VM1 "sudo systemctl stop $service" 2>/dev/null || true
                    ssh $VM2 "sudo systemctl stop $service" 2>/dev/null || true
                    ;;
            esac
        done
        
        sleep 5
        log "Force kill completed"
    fi
fi

# Clean up temporary files
log "Cleaning up temporary files..."

# Clean up log files (optional)
read -p "Do you want to clean up log files? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    ssh $VM3 "rm -f gnb.log*" 2>/dev/null || true
    ssh $VM4 "rm -f ue.log*" 2>/dev/null || true
    log "Log files cleaned up"
fi

# Network cleanup
log "Cleaning up network connections..."

# Clear any remaining SCTP associations
ssh $VM1 "ss -tan | grep 38412 | awk '{print \$4}' | cut -d: -f1 | xargs -I {} ss -K dst {}" 2>/dev/null || true
ssh $VM2 "ss -tan | grep 38412 | awk '{print \$4}' | cut -d: -f1 | xargs -I {} ss -K dst {}" 2>/dev/null || true
ssh $VM3 "ss -tan | grep 38412 | awk '{print \$4}' | cut -d: -f1 | xargs -I {} ss -K dst {}" 2>/dev/null || true

log "5GC testbed shutdown completed successfully!"

# Summary
echo
log "Shutdown Summary:"
log "=================="
log "✓ User Equipment stopped"
log "✓ Radio Access Network stopped"
log "✓ User Plane Function stopped"
log "✓ Core Network Services stopped"
if [[ $REPLY =~ ^[Yy]$ ]]; then
    log "✓ MongoDB stopped"
fi
log "✓ Network connections cleaned up"
echo
log "The 5GC testbed is now ready for the next startup cycle." 